'''
Plot the graph of multiple fields/indicators of a specific stock of specific time span.
matplotlib and pylab libraries are needed for this class.
Details and examples on matplotlib can be referenced at: http://matplotlib.sourceforge.net/gallery.html
 
Created on May, 2012
@author: bruce
'''

# [issue] The raw data does not account for the buffer for indicator period (e.g. moving average). 
# We have to decrease wanted date by a period interval.
from datetime import date
from matplotlib import pyplot as plt
from matplotlib.dates import date2num, num2date
from pylab import legend

from histdata import HistdataBySymbol 


class HistDataPlot:
    def __init__(self, csv_dir="input\\"):  
        self.__csvdir = csv_dir
    
    """
    Plot a time-series of a symbol with properties/indicators and time span that we want.
    ind_pers: indicator - period pairs
    """      
    def plot(self, symbol, ind_pers=[], start_date=None, end_date=None):        
        hist_data_process  = HistdataBySymbol(symbol, self.__csvdir)
        hist_data_process.get_histdailyraws(start_date, end_date)
        # Calculate required indicators.
        hist_data_process.calculate_indicators(ind_pers)
        fig = plt.figure()
        graph = fig.add_subplot(111)
        line_spec = { 'o'     :   'r-',   # red
                      'c'     :   'c-',   # cyan                  
                      'h'     :   'k-',   # black
                      'lo'    :   'y-',   # yellow
                      'adj_c' :   'g-',   # green
                      'ma'    :   'm-',   # magenta
                      'bb'    :   'b-',   # blue
                      'std'   :   'k--', 
                      'ema'   :   'b--', 
                      'wpr'   :   'y--', 
                      'macd'  :   'm--', 
                      'env'   :   'c--', 
                      'rsi'   :   'r--', 
                      'roc'   :   'g:', 
                      'mo'    :   'm:'   
                      }                      
        # date as the x axis            
        t = [date2num(quote.dt) for quote in hist_data_process._histdata]
        # help method to get corresponding property/indicator
        def _get_quotevalue(quote, ind, period):
            val = vars(quote)[ind]
            if type(val) == type({}):
                val = val[period]
            return val
        # draw all the indicators from the list against the time
        for ind, period in ind_pers:
            values = [_get_quotevalue(quote, ind, period) for quote in hist_data_process._histdata]
            graph.plot(t, values, line_spec[ind], label = ind)
            legend(loc = 0)
        # add the time stamp of x-axis for 5 intervals
        start_dt = hist_data_process._histdata[0].dt
        end_dt = hist_data_process._histdata[-1].dt
        startdt_num = int(date2num(start_dt))
        enddt_num = int(date2num(end_dt))
        interval = max(1, int((enddt_num - startdt_num) / 5))
        datenums = range(startdt_num, startdt_num + interval * 5 + 1, interval)
        plot_dates = [self.__format_date_str(num2date(datenum)) for datenum in datenums]
        graph.set_title(symbol)
        graph.set_xticks(datenums)
        graph.set_xticklabels(plot_dates)        
        # rotates and right aligns the x lables.
        fig.autofmt_xdate()
        graph.grid(True)
        plt.show()
    
    def __format_date_str(self, dt_time):
        return str(dt_time.year) + '/' + str(dt_time.month) + '/' + str(dt_time.day)
    
            
if __name__ == "__main__": 
    dataplot = HistDataPlot()
    dataplot.plot('NOK', [('c', 0), ('ma', 10), ('ema', 30), 
                          ('adj_c', 0), ('lo', 0), ('h', 0), ('o', 0), ('env', 16)], 
                  date(2012, 1, 1), date(2012, 5, 1))
    #dataplot.plot('znga', [('c', 0), ('ma', 10), ('ema', 30)])
    #dataplot.plot('goog', [('c', 0), ('ma', 10), ('ema', 30)], date(2012, 1, 1), date(2012, 5, 1))
    #dataplot.plot('aapl', [('c', 0), ('ma', 10), ('ema', 30)], date(2012, 1, 1), date(2012, 5, 1))
    #dataplot.plot('msft', [('c', 0), ('ma', 10), ('ema', 30)], date(2012, 1, 1), date(2012, 5, 1))